///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file psi3_interface.cpp
///
/// \brief   contains definitions of PSI3 access functions
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/blas_lapack.h>
#include <src/psi3_interface.h>

// include headers needed for PSI3
extern "C" {
#include <stdio.h>
#include <stdlib.h>
#include <libciomr/libciomr.h>
#include <psifiles.h>
#include <libchkpt/chkpt.h>
#include <libpsio/psio.h>
#include <libiwl/iwl.h>
#include <libipv1/ip_lib.h>
#include <libqt/qt.h>
#include <math.h>
}

// define global variables for PSI3 libraries
extern "C" {
  char *psi_file_prefix;
  FILE *infile, *outfile;
  char *gprgid() { char *prgid = ":PSICHEM"; return(prgid); }
}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   reads the one and two electron integrals from PSI3
///
/// \param[in]   userinp  user input options
/// \param[out]  enuc     the nuclear repulsion energy
/// \param[out]  oei      the one electron integrals, size = na*na
/// \param[out]  tei      the two electron integrals, size = na*(na+1) * ( na*(na+1) + 2 ) / 8
///
///////////////////////////////////////////////////////////////////////////////////////////////////

void pcps::get_psi3_ints(const pcps::Input & userinp, double & enuc, double * const oei, double * const tei) {

  // initialize psi3
  const int argc = 0;
  char ** argv;
  const int errcod = psi_start(argc, argv, 0);
  if ( errcod != PSI_RETURN_SUCCESS )
    throw pcps::Exception( (boost::format("psi_start failed with code %i in pcps::get_psi3_ints") % errcod).str() );
  //ffile(&infile, "psichem.inp", 2);
  //ffile(&outfile, "psichem.out", 1);
  ip_cwk_add(gprgid());
  tstart(outfile);
  psio_init();

  // open the chkpt file
  chkpt_init(PSIO_OPEN_OLD);

  // load the nuclear repulsion energy
  enuc = chkpt_rd_enuc();

  // get numbers of orbitals
  const int no = chkpt_rd_nmo();       // total
  const int nc = userinp.n_cor_orb();  // core
  const int na = userinp.n_act_orb();  // active

  // check that the numbers of orbitals are sane
  if ( nc + na > no )
    throw pcps::Exception("too many core and/or active orbitals.  nc = %i, na = %i, no = %i") % nc % na % no;

  // check that the number of active orbitals matches the number of lattice sites
  if ( userinp.nsites() != 2*na )
    throw pcps::Exception("number of active orbitals must equal lattice sites:  na = %i, nsites = %i") % na % (userinp.nsites()/2);

  // clost the chkpt files
  chkpt_close();

  // create temporary vectors into which we will load the integrals
  std::vector<double> temp_int1(no*no, 0.0);
  std::vector<double> temp_int2( no*(no+1)/2 * ( no*(no+1)/2 + 1 ) / 2, 0.0);

  // get the number of elements in the one electron integral's lower triangle
  const int ntri = no*(no+1)/2;

  // create a vector used in loading one electron data from PSI3
  std::vector<double> lower_triangle(ntri, 0.0);

  // load the one electron integrals in lower triangle format
  iwl_rdone(PSIF_OEI, PSIF_MO_OEI, &lower_triangle.at(0), ntri, 0, 0, outfile);

  // expand the one electron integrals into full-matrix storage
  pcps::unpack_psi3_lower_triangle(no, &lower_triangle.at(0), &temp_int1.at(0));

  // precompute offset array used in loading two electron integrals
  std::vector<int> ioff(ntri*ntri, 0);
  for (int i = 0; i < ioff.size(); i++)
    ioff[i] = i*(i+1)/2;

  // load two electron integrals
  iwl_rdtwo(PSIF_MO_TEI, &temp_int2.at(0), &ioff.at(0), no, 0, 0, 0, outfile);

  // finalize psi3
  psio_done();
  psi_stop();

  // extract active space one electron integrals
  for (int p = 0; p < na; p++)
  for (int q = 0; q < na; q++) {
    oei[na*p+q] = temp_int1[no*(p+nc)+(q+nc)];
    for (int i = 0; i < nc; i++)
      oei[na*p+q] +=   2.0 * temp_int2[pcps::Hamiltonian<double>::tei_compound_index(i, i, p+nc, q+nc)]
                     - 1.0 * temp_int2[pcps::Hamiltonian<double>::tei_compound_index(i, p+nc, i, q+nc)];
  }

  // extract active space two electron integrals
  for (int p = 0; p < na; p++)
  for (int q = 0; q <= p; q++)
  for (int r = 0; r < na; r++)
  for (int s = 0; s <= r; s++) {
    const int tmp_cmpd = pcps::Hamiltonian<double>::tei_compound_index(p+nc, q+nc, r+nc, s+nc);
    const int act_cmpd = pcps::Hamiltonian<double>::tei_compound_index(p, q, r, s);
    tei[act_cmpd] = temp_int2[tmp_cmpd];
  }

  // add core orbital contribution to the nuclear energy (nuclear energy is used to also hold core energy)
  for (int i = 0; i < nc; i++)
    enuc += 2.0 * temp_int1[no*i+i] + temp_int2[pcps::Hamiltonian<double>::tei_compound_index(i, i, i, i)];
  for (int i = 0; i < nc; i++)
  for (int j = 0; j <  i; j++)
    enuc +=   4.0 * temp_int2[pcps::Hamiltonian<double>::tei_compound_index(i, i, j, j)]
            - 2.0 * temp_int2[pcps::Hamiltonian<double>::tei_compound_index(i, j, i, j)];

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   throws an error because PSI3 doesn't use complex integrals
///
/// \param[in]   userinp  user input options
/// \param[out]  enuc     the nuclear repulsion energy
/// \param[out]  oei      the one electron integrals, size = na*na
/// \param[out]  tei      the two electron integrals, size = na*(na+1) * ( na*(na+1) + 2 ) / 8
///
///////////////////////////////////////////////////////////////////////////////////////////////////


void pcps::get_psi3_ints(const pcps::Input & userinp, double & enuc, std::complex<double> * const oei,
                                                                     std::complex<double> * const tei) {

  throw pcps::Exception("Cannot read complex integrals from PSI3.  Please use scalar_type = real");

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief   expands the psi3 lower triangle storage for one electron integrals
///
///   This function adds the one electron integral data stored in the lower triangle to the one
///   electron integral data stored in the full matrix.
///
/// \param[in]      n        the number of spatial orbitals
/// \param[in]      lt       the one electron integrals in lower triangle storage, size n*(n+1)/2
/// \param[in,out]  oei      the one electron integrals in full-matrix storage, size = n*n
///
///////////////////////////////////////////////////////////////////////////////////////////////////

void pcps::unpack_psi3_lower_triangle(const int n, const double * const lt, double * const oei) {

  for (int i = 0; i < n; i++) {

    // diagonal elements
    oei[i*n+i] += lt[ i*(i+1)/2 + i ];

    // off diagonal elements
    for (int j = 0; j < i; j++) {
      oei[i*n+j] += lt[ i*(i+1)/2 + j ];
      oei[j*n+i] += lt[ i*(i+1)/2 + j ];
    }

  }

}
